425A - Sereja and Swaps - CodeForces Solution

#include <iostream>
#include <map>
#include <vector>

template <typename T>
std::istream& operator >>(std::istream& input, std::vector<T>& v)
{
    for (T& a : v)
        input >> a;

    return input;
}

void answer(int x)
{
    std::cout << x << '\n';
}

void consume(std::map<int, unsigned>& f, int x, unsigned d)
{
    const auto it = f.find(x);
    it->second -= d;
    if (it->second == 0)
        f.erase(it);
}

void solve(const std::vector<int>& a, size_t k)
{
    const size_t n = a.size();

    std::map<int, unsigned> f;
    for (const int x : a)
        ++f[x];

    const auto check = [&](std::map<int, unsigned> s) {
        int v = 0;

        std::map<int, unsigned> t = f;
        for (const auto& e : s) {
            v += e.first * int(e.second);
            consume(t, e.first, e.second);
        }
        for (unsigned q = 0; q < k && !s.empty() && !t.empty() && s.begin()->first < t.rbegin()->first; ++q) {
            v -= s.begin()->first;
            consume(s, s.begin()->first, 1);

            v += t.rbegin()->first;
            consume(t, t.rbegin()->first, 1);
        }

        return v;
    };

    int x = a[0];
    for (size_t i = 0; i < n; ++i) {
        std::map<int, unsigned> s;
        for (size_t j = i; j < n; ++j) {
            ++s[a[j]];
            x = std::max(x, check(s));
        }
    }

    answer(x);
}

int main()
{
    std::cin.tie(nullptr)->sync_with_stdio(false);

    size_t n, k;
    std::cin >> n >> k;

    std::vector<int> a(n);
    std::cin >> a;

    solve(a, k);

    return 0;
}